# encoding=utf-8

"""IPredictHandler
    - 定义主程序中各个模块的处理通用接口
"""
from seg_common.BatchLoader import BatchLoaderFactory
from seg_common.PredictService import PredictFactory
from seg_common.BatchLoader import BatchCommonDataLoader
from seg_common.logging import ConsoleService
from seg_system import ApplicationConfig
from seg_system.LongTaskStatusUpload import IStatusUpload
from seg_system.LongTaskStatusUpload import MultiStatusUploadImpl


class IPredictHandler:
    @staticmethod
    def predict(from_path: list, to_path: list,
                batch_factory: BatchLoaderFactory,
                predict_service: PredictFactory,
                batch_size: int = ApplicationConfig.SystemConfig.BATCH_SIZE_NORMAL,
                pre_message: str = "", **kwargs):

        # get kwargs
        tmp_list = kwargs.get('tmp_list', None)
        ceil_list = kwargs.get('ceil_list', None)

        for each_from, each_to in zip(from_path, to_path):
            batch_factory.load_data(each_from)
            data_loader = BatchCommonDataLoader.make_data_loader(batch_factory, batch_size=batch_size)
            for _, (tensor, index) in enumerate(data_loader):
                o = predict_service.predict(tensor)
                token = batch_factory.get_tokens(list(index.cpu().numpy()))

                # build kwargs
                d = {'each_from': each_from}
                index = from_path.index(each_from)
                if tmp_list:
                    d["tmp_path"] = tmp_list[index]
                if ceil_list:
                    d['ceil_path'] = ceil_list[index]

                predict_service.save(o, each_to, token, **d)
                ConsoleService.console_log("{}\n\tpath: {}\n\ttoken: {}".format(pre_message, each_from, str(token)))
                IStatusUpload.update_state_from_kwargs(**kwargs, token=token)
